from ..config import OpPrior, ReducedFraction, Space
from .base import Op1, Op2, OpConst, OpF1
from .function import OpLn


class OpNeg(Op1):
    op = disp = '-'
    lamda = lambda _, r: -r
    level = OpPrior.neg
    addSpace = Space.neg

    def diff(self, dx):
        return -self.r.diff(dx)

    def simple(self):
        r = self.r.simple()
        if isinstance(r, OpConst):
            return OpConst(self.lamda(r.value))
        elif isinstance(r, OpNeg):
            return r.r
        return self.same(r)


class OpId(OpF1):
    op = "id"
    disp = "Id"
    lamda = lambda _, r: r

    def diff(self, dx):
        return self.r.diff(dx)


class OpAdd(Op2):
    op = disp = '+'
    lamda = lambda _, l, r: l + r
    level = OpPrior.add
    addSpace = Space.add

    def diff(self, dx):
        return self.l.diff(dx) + self.r.diff(dx)

    def simple(self):
        l, r = self.l.simple(), self.r.simple()
        if isinstance(l, OpConst):
            if l == 0:
                return r
            elif isinstance(r, OpConst):
                return OpConst(self.lamda(l.value, r.value))
        if isinstance(r, OpConst) and r == 0:
            return l
        if isinstance(r, OpNeg):
            return l - r.r
        return self.same(l, r)


class OpSub(Op2):
    op = disp = '-'
    lamda = lambda _, l, r: l - r
    level = OpPrior.sub
    addSpace = Space.sub

    def diff(self, dx):
        return self.l.diff(dx) - self.r.diff(dx)

    def simple(self):
        l, r = self.l.simple(), self.r.simple()
        if isinstance(l, OpConst):
            if l == 0:
                return -r
            elif isinstance(r, OpConst):
                return OpConst(self.lamda(l.value, r.value))
        if isinstance(r, OpConst) and r == 0:
            return l
        if isinstance(r, OpNeg):
            return l + r.r
        if l == r:
            return OpConst(0)
        return self.same(l, r)


class OpMul(Op2):
    op = '*'
    disp = '·'
    lamda = lambda _, l, r: l * r
    level = OpPrior.mul
    addSpace = Space.mul

    def diff(self, dx):
        return self.l.diff(dx) * self.r + self.l * self.r.diff(dx)

    def simple(self):
        l, r = self.l.simple(), self.r.simple()
        if isinstance(l, OpConst):
            if l == 0:
                return OpConst(0)
            elif l == 1:
                return r
            elif isinstance(r, OpConst):
                return OpConst(self.lamda(l.value, r.value))
        if isinstance(r, OpConst):
            if r == 0:
                return OpConst(0)
            elif r == 1:
                return l
            else:
                l, r = r, l
        return self.same(l, r)


class OpDiv(Op2):
    op = disp = '/'
    lamda = lambda _, l, r: l / r
    level = OpPrior.div
    addSpace = Space.div

    def diff(self, dx):
        return (self.l.diff(dx) * self.r - self.l * self.r.diff(dx)) / (self.r * self.r)

    def simple(self):
        l, r = self.l.simple(), self.r.simple()
        if isinstance(l, OpConst):
            if l == 0:
                return OpConst(0)
            elif isinstance(r, OpConst) and ReducedFraction:
                return OpConst(self.lamda(l.value, r.value))
        if isinstance(r, OpConst):
            if r == 0:
                return OpConst(float("inf"))
            elif r == 1:
                return l
        if l == r:
            return OpConst(1)
        return self.same(l, r)


class OpPow(Op2):
    op = "**"
    disp = '^'
    lamda = lambda _, l, r: l ** r
    level = OpPrior.pow
    addSpace = Space.pow

    def diff(self, dx):
        if not (difL := self.l.diff(dx)):
            return self.l ** self.r * OpLn(self.l) * self.r.diff(dx)
        if self.r.diff(dx):
            raise NotImplementedError()
        else:
            return self.r * self.l ** (self.r - 1) * difL

    def simple(self):
        l, r = self.l.simple(), self.r.simple()
        if isinstance(l, OpConst):
            if l == 0:
                return OpConst(0)
            elif l == 1:
                return OpConst(1)
            elif isinstance(r, OpConst):
                return OpConst(self.lamda(l.value, r.value))
        if isinstance(r, OpConst):
            if r == 0:
                return OpConst(1)
            elif r == 1:
                return l
        return self.same(l, r)
